package SSF.OS.NetFlow;

import com.renesys.raceway.DML.*;
import java.util.*;
import SSF.OS.*;
import SSF.OS.TCP.*;
import SSF.OS.NetFlow.Util.*;
import SSF.Net.*;
import java.util.*;
import SSF.OS.Timer;

/**
 * IpFlowCollector With Domain-support
 */
public class IpFlowCollectorWD extends IpFlowCollector
{
    public int level = 1;
    
    /** table to store NIC type information*/
    protected Hashtable nicTable;
    protected int nicTableSize=5;

    /** 
     * All options as in IpFlowCollector. Plus:
     * <UL>
     * <LI> level -- how many "level" of nhi addresses is used to represent a domain.
     * <LI> nic_table_size -- initial size of the NIC table.
     * </UL>
     */
    public void config(IP ipSession, Configuration cfg) throws configException
    {
	super.config(ipSession, cfg);
	String vals = (String)cfg.findSingle("level");
	if (vals != null)
	    level = Integer.valueOf(vals).intValue();
	vals = (String)cfg.findSingle("nic_table_size");
	if (vals !=null)
	    nicTableSize = Integer.valueOf(vals).intValue();
    }

    /**
     * It assumes that the host doesn't have a large number of NIC attacked.
     * Otherwise, the performance will not be good because of the small size
     * of the nicTable..
     */
    public void init()
    {
	super.init();

	//get the link information about the NIC and put it the table.
	nicTable = new Hashtable(5);
	NIC curNIC = null;

	for (int i=0; i<ipSess.INTERFACE_COUNT; i++) {
	    curNIC = ipSess.INTERFACE_SET[i];
	    if (curNIC.link_hw instanceof ptpLinkLayer) {
		//compare the nhi prefix of the nic with its peer's
		if (NHIDemuxer.getNhiPrefix
		    (((Host)(((ptpLinkLayer)(curNIC.link_hw)).peer(curNIC)
			    .inGraph())).nhi, 
		     level)
		    .equals(
		    NHIDemuxer.getNhiPrefix(((Host)(ipSess.inGraph())).nhi, 
					       level))) {
		    //internal NIC
		    nicTable.put(new Integer(curNIC.ID), 
				 new Byte(IpNetFlowWD.INTERNAL));
		} else {
		    // not the same prefixes
		    nicTable.put(new Integer(curNIC.ID),
				 new Byte(IpNetFlowWD.EXTERNAL));
		}
	    } else {
		// not ptp link
		nicTable.put(new Integer(curNIC.ID),
			     new Byte(IpNetFlowWD.NOT_PTP));
	    }
	}
    }

    /**
     * Same as in IpFlowCollector except it uses IpNetFlowWD.
     */
    public void receive(ProtocolMessage packet, ProtocolSession fromSession,
			ProtocolSession toSession)
    {
	IpHeader ipPacket = (IpHeader)packet;
	IpNetFlowWD flow = null;
	IpNetFlowWD flowList = null;
	Long key;

	//-- check whether it's the protocol interested in.
	//if (protocol_type != ipPacket.PROTOCOL_NO)
	//  return;         //-- not the protocol wanted, pass it.

	if (protocol_type != ALL_PROTOCOLS) {
	    if (ipPacket.PROTOCOL_NO == Protocols.TCP_PRTL_NUM) {
		
		if (protocol_type != Protocols.TCP_PRTL_NUM
		    && protocol_type != Protocols.BGP_PRTL_NUM)
		    return;
		
		// TCP or protocols based on TCP, check further
		switch (protocol_type) {
		case Protocols.TCP_PRTL_NUM :
		    break;
		case Protocols.BGP_PRTL_NUM:
		    //check whether it's BGP, if it is, go on.
		    Object payLoad = ipPacket.payload().payload();
		    if (payLoad == null) return;
		    if (!protocolMsgTypeClass.isInstance(payLoad))
			return;
		    break;
		case Protocols.APP_PRTL_NUM:
		    //further check needed to make sure it is interested.
		    break;
		default:
		    return;      //not the ones that are interested.
		}
	    }
	    //if it is "APP_PRTL_NUM", it will pass without further change. !!!!
	    //"APP_PRTL_NUM" is for Applications, that is, everthing.
	    else if (ipPacket.PROTOCOL_NO != Protocols.APP_PRTL_NUM) {
		//in current implementation
		//no other protocols are based on either UDP or sOSPF use
		if (ipPacket.PROTOCOL_NO != protocol_type)
		    return;
	    }
	}

	//-- check whether it's from NIC, if not return.
	short input= LOCAL;
	short output= LOCAL;

	if (nicClass.isInstance(fromSession))
	    input = (short)(((NIC)fromSession).ID);
	if (nicClass.isInstance(toSession))
	    output = (short)(((NIC)toSession).ID);

	float curTime = getSimTime();

	//-- compute the key for this ipPacket
	key = IpNetFlow.makeKey(ipPacket);

	//-- check whether the packet belongs to an existing flow
	flowList = (IpNetFlowWD)conOpenTable.get(key);

	if (flowList == null) {
	    //-- new flow, add to this entry
	    //System.err.println("No entry, add new Flow");
	    flow = new IpNetFlowWD(ipPacket, input,
				   (input==LOCAL?IpNetFlowWD.INTERNAL:
				    ((Byte)nicTable.get(new Integer(input)))
				   .byteValue()),
				   output, 
				   (output==LOCAL?IpNetFlowWD.INTERNAL:
				    ((Byte)nicTable.get(new Integer(output)))
				   .byteValue()),
				   curTime,
				   ((Host)inGraph()).nhi, srcMask, dstMask);
	    conOpenTable.put(key, flow);
	}
	else {
	    //-- a list that potentially has the flow we want was returned
	    //-- check the elt  one by one, try to find matched flow
	    //-- ! Notice that it also check whether the input / output NIC 
	    //-- is the same.
	    //-- If not, it means it might split into a new flow,
	    //-- or, this is the point where two flows merge,
	    //-- so keep checking the followig elements until you find the one
	    //-- or you create a new one.
	    for (flow = flowList; flow!=null; flow=(IpNetFlowWD)(flow.next)) {
		if (flow.acceptPacket(ipPacket) 
		    && (output == flow.output)
		    && (input==flow.input) 
		    && (curTime - flow.Last <= maxInactiveTime))
		    break;
	    }

	    if (flow == null) {
		//-- no match, add a new flow to this list
		//System.err.println("add New Flow");
		flow=new IpNetFlowWD(ipPacket, input,
				     (input==LOCAL?IpNetFlowWD.INTERNAL:
				      ((Byte)nicTable.get(new Integer(input)))
				     .byteValue()),
				     output, 
				     (output==LOCAL?IpNetFlowWD.INTERNAL:
				      ((Byte)nicTable.get(new Integer(output)))
				      .byteValue()),
				     curTime,
				     ((Host)inGraph()).nhi, srcMask, dstMask);
		conOpenTable.put(key, flow);
		
		flow.next = flowList;
		conOpenTable.put(key, flow);
	    }
	    else {                //-- there is a match, just update
		//for debug
		//System.err.println("Update");
		flow.update(packet, curTime);
	    }
	}

	//-- check whether this is a TCP FIN packet
	if (protocol_type == Protocols.TCP_PRTL_NUM) {	    
	    //-- check whether this is a FIN packet, 	
	    if (((TCP_Header)(ipPacket.payload())).flags 
		== TCP_Header.FIN) {
		//-- if yes, it's the end of this TCP flow
		//-- move it to the connectoin finish List
		//??? I can do slightly more efficient here by NOT calling
		//??? the removeFlow function... the flowList and the flow
		//??? are both known by now.
		conOpenTable.removeFlow(key, flow);
		//conFinishList.addHead(flow);
		conFinishList.addLast(flow);
	    }
	}
    }

}
